#include "kernel_operator.h"
using namespace AscendC;

class KernelNLLLossMean {
public:
    __aicore__ inline KernelNLLLossMean() {}
    __aicore__ inline void Init(TPipe* pipe, 
                            GM_ADDR xAddr,GM_ADDR targetAddr,GM_ADDR weightAddr,GM_ADDR yAddr,GM_ADDR workAddr,
                            const uint32_t& N, const uint32_t& C, const uint32_t& bn, const uint32_t& fbn
                            , const uint32_t& fn
                            )
    {
        this->N = N;
        this->C = C;
        this->bn = bn;
        this->fbn = fbn;
        this->fn = fn;
        xGm.SetGlobalBuffer((__gm__ float*)xAddr, N*C);
        targetGm.SetGlobalBuffer((__gm__ int32_t*)targetAddr, N);
        weightGm.SetGlobalBuffer((__gm__ float*)weightAddr, C);
        yGm.SetGlobalBuffer((__gm__ float*)yAddr, 1);

    }

    __aicore__ inline void Process()
    {
        float sum=0.0f, weightSum=0.0f;
        for(int n = 0; n < N; ++n){
            int tn = targetGm.GetValue(n);
            float wtn = weightGm.GetValue(tn);
            float xtn = xGm.GetValue(n*C+tn);
            sum += -wtn*xtn;
            weightSum += wtn;
        }
        yGm.SetValue(0, sum/weightSum);
    }
private:
    TPipe* pipe;
    GlobalTensor<float> xGm;
    GlobalTensor<int32_t> targetGm;
    GlobalTensor<float> weightGm;
    GlobalTensor<float> yGm;

    uint32_t N;
    uint32_t C;
    uint32_t bn;
    uint32_t fbn;
    uint32_t fn;

};
